#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import asyncio
import fractions
import io
import logging
import time
from typing import Optional

import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from aiortc import VideoStreamTrack
from av import VideoFrame

logger = logging.getLogger(__name__)

class VideoGenerator:
    """视频生成器 - 循环播放JPEG图片生成视频流"""
    
    def __init__(self, width: int = 640, height: int = 480, fps: int = 30):
        self.width = width
        self.height = height
        self.fps = fps
        self.frame_duration = 1.0 / fps
        
        self.images = []
        self.current_frame_index = 0
        self.is_running = False
        
        # 生成示例图片
        self._generate_sample_images()
        
        logger.info(f"视频生成器初始化: {width}x{height}@{fps}fps, {len(self.images)}张图片")
    
    def _generate_sample_images(self):
        """生成示例JPEG图片"""
        colors = [
            (255, 100, 100),  # 红色
            (100, 255, 100),  # 绿色
            (100, 100, 255),  # 蓝色
            (255, 255, 100),  # 黄色
            (255, 100, 255),  # 紫色
            (100, 255, 255),  # 青色
        ]
        
        try:
            # 尝试加载系统字体
            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 40)
        except:
            try:
                font = ImageFont.truetype("arial.ttf", 40)
            except:
                font = ImageFont.load_default()
        
        for i, color in enumerate(colors):
            # 创建PIL图像
            img = Image.new('RGB', (self.width, self.height), color)
            draw = ImageDraw.Draw(img)
            
            # 添加文字
            text = f"Frame {i+1}"
            
            # 计算文字位置（居中）
            try:
                bbox = draw.textbbox((0, 0), text, font=font)
                text_width = bbox[2] - bbox[0]
                text_height = bbox[3] - bbox[1]
            except:
                # 兼容旧版本PIL
                text_width, text_height = draw.textsize(text, font=font)
            
            x = (self.width - text_width) // 2
            y = (self.height - text_height) // 2
            
            # 绘制文字（黑色背景，白色文字）
            draw.rectangle([x-10, y-10, x+text_width+10, y+text_height+10], fill=(0, 0, 0))
            draw.text((x, y), text, fill=(255, 255, 255), font=font)
            
            # 添加时间戳
            timestamp = time.strftime("%H:%M:%S")
            draw.text((10, 10), timestamp, fill=(255, 255, 255), font=font)
            
            # 添加一些装饰图形
            self._add_decorations(draw, i)
            
            # 转换为numpy数组
            img_array = np.array(img)
            
            # 转换为BGR格式（OpenCV格式）
            img_bgr = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)
            
            self.images.append(img_bgr)
            
            logger.debug(f"生成图片 {i+1}: {color}")
    
    def _add_decorations(self, draw, frame_index):
        """添加装饰图形"""
        # 绘制一些圆形
        for i in range(3):
            x = 50 + i * 100 + frame_index * 20
            y = self.height - 100
            radius = 20 + frame_index * 2
            
            # 确保圆形在画布内
            x = x % (self.width - radius * 2)
            
            draw.ellipse([x, y, x + radius * 2, y + radius * 2], 
                        fill=(255, 255, 255), outline=(0, 0, 0), width=2)
        
        # 绘制进度条
        progress_width = int((self.width - 40) * (frame_index + 1) / 6)
        draw.rectangle([20, self.height - 30, 20 + progress_width, self.height - 10], 
                      fill=(255, 255, 255))
        draw.rectangle([20, self.height - 30, self.width - 20, self.height - 10], 
                      outline=(255, 255, 255), width=2)
    
    async def start(self):
        """启动视频生成器"""
        self.is_running = True
        logger.info("视频生成器已启动")
    
    async def stop(self):
        """停止视频生成器"""
        self.is_running = False
        logger.info("视频生成器已停止")
    
    def get_next_frame(self) -> np.ndarray:
        """获取下一帧图像"""
        if not self.images:
            # 如果没有图像，返回黑色图像
            return np.zeros((self.height, self.width, 3), dtype=np.uint8)
        
        # 获取当前帧
        frame = self.images[self.current_frame_index].copy()
        
        # 更新帧索引
        self.current_frame_index = (self.current_frame_index + 1) % len(self.images)
        
        return frame
    
    def get_video_track(self) -> 'GeneratedVideoTrack':
        """获取视频轨道"""
        return GeneratedVideoTrack(self)

class GeneratedVideoTrack(VideoStreamTrack):
    """生成的视频轨道"""
    
    def __init__(self, generator: VideoGenerator):
        super().__init__()
        self.kind = "video"  # 明确设置轨道类型
        self.generator = generator
        self.start_time = time.time()
        
        logger.info("视频轨道已创建")
    
    async def recv(self) -> VideoFrame:
        """接收视频帧"""
        # 计算当前应该的帧时间
        current_time = time.time()
        elapsed_time = current_time - self.start_time
        expected_frame = int(elapsed_time * self.generator.fps)
        
        # 等待到正确的时间
        next_frame_time = self.start_time + (expected_frame + 1) / self.generator.fps
        wait_time = next_frame_time - current_time
        
        if wait_time > 0:
            await asyncio.sleep(wait_time)
        
        # 获取图像帧
        img = self.generator.get_next_frame()
        
        # 转换为VideoFrame
        frame = VideoFrame.from_ndarray(img, format="bgr24")
        frame.pts = expected_frame
        frame.time_base = fractions.Fraction(1, self.generator.fps)
        
        # 添加调试日志（每30帧记录一次，避免日志过多）
        if expected_frame % 30 == 0:
            logger.info(f"发送视频帧: {expected_frame}, 尺寸: {img.shape}")
        
        return frame

class StaticImageGenerator(VideoGenerator):
    """静态图像生成器 - 用于测试"""
    
    def __init__(self, image_path: str, width: int = 640, height: int = 480, fps: int = 30):
        self.width = width
        self.height = height
        self.fps = fps
        self.frame_duration = 1.0 / fps
        
        self.images = []
        self.current_frame_index = 0
        self.is_running = False
        
        # 加载静态图像
        self._load_static_image(image_path)
        
        logger.info(f"静态图像生成器初始化: {image_path}")
    
    def _load_static_image(self, image_path: str):
        """加载静态图像"""
        try:
            # 使用OpenCV加载图像
            img = cv2.imread(image_path)
            if img is None:
                raise ValueError(f"无法加载图像: {image_path}")
            
            # 调整图像大小
            img = cv2.resize(img, (self.width, self.height))
            
            # 创建多个副本（添加不同的时间戳）
            for i in range(6):
                img_copy = img.copy()
                
                # 添加时间戳
                timestamp = f"Frame {i+1} - {time.strftime('%H:%M:%S')}"
                cv2.putText(img_copy, timestamp, (10, 30), 
                           cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
                
                self.images.append(img_copy)
                
        except Exception as e:
            logger.error(f"加载静态图像失败: {e}")
            # 回退到生成示例图片
            self._generate_sample_images()

if __name__ == '__main__':
    # 测试视频生成器
    async def test_generator():
        generator = VideoGenerator()
        await generator.start()
        
        # 生成几帧进行测试
        for i in range(10):
            frame = generator.get_next_frame()
            print(f"生成帧 {i+1}: {frame.shape}")
            await asyncio.sleep(0.1)
        
        await generator.stop()
    
    asyncio.run(test_generator())